from enum import Enum
from typing import Any, Literal, TypeAlias, Union
import pandas
from pydantic import BaseModel, Field
from ..dataclass.geoProfiles import MaterialTable
from ..dataclass.tables import TableData
from ..pipeline.pipeData import ArrayAttributeSchema, BooleanAttributeSchema, ObjectAttributeSchema, StringAttributeSchema

CleaningRule: TypeAlias = Union[NaNCheckRule, DataTypeRule, ValueRangeRule, DuplicateRowsRule, ConsistencyRule, CustomValidationRule]

def format_dataframe_type(df: Any, dtype: dict[str, str], to_numeric: bool = True) -> Any:
    """
    Convert the dataframe to the specific type for each column.
    
    Args:
        df: The dataframe to convert.
        dtype: The type to convert the dataframe to. The key is the column name, the value is the type.
               Keys that don't exist in the dataframe will be ignored.
        to_numeric: If it's True, invalid string like `""` will be convert to `NaN` when converting to numeric type.
    
    Returns:
        The converted dataframe of the same type as the input (DataFrame or TableData).
    """
    ...
def create_cleaning_rule_ui_schema() -> ObjectAttributeSchema:
    """Create UI schema for cleaning rule configuration."""
    ...
def create_cleaning_config_ui_schema() -> ObjectAttributeSchema:
    """Create UI schema for cleaning module configuration."""
    ...

class CleaningActionType(str, Enum):
    """Actions that can be taken when data cleaning rules are triggered."""
    IGNORE: Any
    WARNING: Any
    GDI_WARNING: Any
    ERROR: Any
    REMOVE_ROWS: Any
    FILL_VALUE: Any
    DROP_COLUMNS: Any
    CONVERT_TYPE: Any
    SKIP_VALIDATION: Any

class CleaningSeverity(str, Enum):
    """Severity levels for data quality issues."""
    LOW: Any
    MEDIUM: Any
    HIGH: Any
    CRITICAL: Any

class BaseCleaningRule(BaseModel):
    """Base model for all data cleaning rules."""
    rule_name: str
    description: str | None
    enabled: bool
    action: CleaningActionType
    severity: CleaningSeverity
    message_template: str | None
    stop_on_error: bool
    target_tables: list[str] | None
    target_columns: list[str] | None
    exclude_tables: list[str] | None
    exclude_columns: list[str] | None

class NaNCheckRule(BaseCleaningRule):
    """Rule for checking missing/NaN values in data."""
    rule_type: Literal[Any]
    threshold_ratio: float
    threshold_count: int | None
    fill_value: Any
    fill_method: Literal[Any, Any, Any, Any, Any, Any, Any] | None

class DataTypeRule(BaseCleaningRule):
    """Rule for validating and converting data types."""
    rule_type: Literal[Any]
    expected_types: dict[str, str]
    convert_types: bool
    allow_null: bool
    strict_mode: bool

class ValueRangeRule(BaseCleaningRule):
    """Rule for validating value ranges (numeric data)."""
    rule_type: Literal[Any]
    column_ranges: dict[str, Any]
    inclusive_bounds: bool
    handle_outliers: Literal[Any, Any, Any]

class DuplicateRowsRule(BaseCleaningRule):
    """Rule for detecting and handling duplicate rows."""
    rule_type: Literal[Any]
    subset_columns: list[str] | None
    keep_first: bool
    ignore_index: bool

class ConsistencyRule(BaseCleaningRule):
    """Rule for checking data consistency across columns or tables."""
    rule_type: Literal[Any]
    consistency_checks: list[Any]
    cross_table_checks: bool

class CustomValidationRule(BaseCleaningRule):
    """Rule for custom validation using user-defined functions."""
    rule_type: Literal[Any]
    validation_function: str
    function_parameters: dict[str, Any]

class CleaningConfiguration(BaseModel):
    """Configuration for a data cleaning module."""
    module_name: str
    rules: list[CleaningRule]
    strict_mode: bool
    parallel_execution: bool
    max_errors: int | None
    generate_report: bool
    report_level: Literal[Any, Any, Any]
    save_backup: bool
    chunk_size: int | None
    memory_limit_mb: int | None

class CleaningRuleTemplate(BaseModel):
    """Template for creating commonly used cleaning rule configurations."""
    template_name: str
    template_description: str | None
    rule_templates: list[Any]
    suggested_order: list[str] | None
    def create_basic_cleaning_template(cls) -> Any:
        """Create a basic cleaning template with common rules."""
        ...
    def create_strict_validation_template(cls) -> Any:
        """Create a strict validation template."""
        ...
